# --- browser_utils/model_management.py ---
# 浏览器模型管理相关功能模块

import asyncio
import json
import logging
import os
from typing import Optional

from playwright.async_api import (
    Page as AsyncPage,
)
from playwright.async_api import (
    expect as expect_async,
)

# 导入配置和模型
from config import (
    AI_STUDIO_URL_PATTERN,
    INPUT_SELECTOR,
)
from logging_utils import set_request_id

logger = logging.getLogger("AIStudioProxyServer")

# ==================== 强制UI状态设置功能 ====================


async def _verify_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> dict:
    """
    验证UI状态设置是否正确

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        dict: 包含验证结果的字典
    """
    set_request_id(req_id)
    try:
        logger.info("验证UI状态设置...")

        # 获取当前localStorage设置
        prefs_str = await page.evaluate(
            "() => localStorage.getItem('aiStudioUserPreference')"
        )

        if not prefs_str:
            logger.warning(" localStorage.aiStudioUserPreference 不存在")
            return {
                "exists": False,
                "isAdvancedOpen": None,
                "areToolsOpen": None,
                "needsUpdate": True,
                "error": "localStorage不存在",
            }

        try:
            prefs = json.loads(prefs_str)
            is_advanced_open = prefs.get("isAdvancedOpen")
            are_tools_open = prefs.get("areToolsOpen")

            # 检查是否需要更新
            needs_update = (is_advanced_open is not True) or (
                are_tools_open is not True
            )

            result = {
                "exists": True,
                "isAdvancedOpen": is_advanced_open,
                "areToolsOpen": are_tools_open,
                "needsUpdate": needs_update,
                "prefs": prefs,
            }

            logger.info(
                f" UI状态验证结果: isAdvancedOpen={is_advanced_open}, areToolsOpen={are_tools_open} (期望: True), needsUpdate={needs_update}"
            )
            return result

        except json.JSONDecodeError as e:
            logger.error(f" 解析localStorage JSON失败: {e}")
            return {
                "exists": False,
                "isAdvancedOpen": None,
                "areToolsOpen": None,
                "needsUpdate": True,
                "error": f"JSON解析失败: {e}",
            }

    except asyncio.CancelledError:
        raise
    except Exception as e:
        logger.error(f" 验证UI状态设置时发生错误: {e}")
        return {
            "exists": False,
            "isAdvancedOpen": None,
            "areToolsOpen": None,
            "needsUpdate": True,
            "error": f"验证失败: {e}",
        }


async def _force_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> bool:
    """
    强制设置UI状态

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        bool: 设置是否成功
    """
    try:
        logger.info(" 开始强制设置UI状态...")

        # 首先验证当前状态
        current_state = await _verify_ui_state_settings(page, req_id)

        if not current_state["needsUpdate"]:
            logger.info(" UI状态已正确设置，无需更新")
            return True

        # 获取现有preferences或创建新的
        prefs = current_state.get("prefs", {})

        # 强制设置关键配置
        prefs["isAdvancedOpen"] = True
        prefs["areToolsOpen"] = True

        # 保存到localStorage
        prefs_str = json.dumps(prefs)
        await page.evaluate(
            "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
            prefs_str,
        )

        logger.info(" 已强制设置: isAdvancedOpen=true, areToolsOpen=true")

        # 验证设置是否成功
        verify_state = await _verify_ui_state_settings(page, req_id)
        if not verify_state["needsUpdate"]:
            logger.info(" UI状态设置验证成功")
            return True
        else:
            logger.warning(" UI状态设置验证失败，可能需要重试")
            return False

    except asyncio.CancelledError:
        raise
    except Exception as e:
        logger.error(f" 强制设置UI状态时发生错误: {e}")
        return False


async def _force_ui_state_with_retry(
    page: AsyncPage,
    req_id: str = "unknown",
    max_retries: int = 3,
    retry_delay: float = 1.0,
) -> bool:
    """
    带重试机制的UI状态强制设置

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志
        max_retries: 最大重试次数
        retry_delay: 重试延迟（秒）

    Returns:
        bool: 设置是否最终成功
    """
    for attempt in range(1, max_retries + 1):
        logger.info(f" 尝试强制设置UI状态 (第 {attempt}/{max_retries} 次)")

        success = await _force_ui_state_settings(page, req_id)
        if success:
            logger.info(f" UI状态设置在第 {attempt} 次尝试中成功")
            return True

        if attempt < max_retries:
            logger.warning(f" 第 {attempt} 次尝试失败，{retry_delay}秒后重试...")
            await asyncio.sleep(retry_delay)
        else:
            logger.error(f" UI状态设置在 {max_retries} 次尝试后仍然失败")

    return False


async def _verify_and_apply_ui_state(page: AsyncPage, req_id: str = "unknown") -> bool:
    """
    验证并应用UI状态设置的完整流程

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        bool: 操作是否成功
    """
    try:
        logger.info(" 开始验证并应用UI状态设置...")

        # 首先验证当前状态
        state = await _verify_ui_state_settings(page, req_id)

        logger.info(
            f" 当前UI状态: exists={state['exists']}, isAdvancedOpen={state['isAdvancedOpen']}, areToolsOpen={state['areToolsOpen']}, needsUpdate={state['needsUpdate']}"
        )

        if state["needsUpdate"]:
            logger.info(" 检测到UI状态需要更新，正在应用强制设置...")
            return await _force_ui_state_with_retry(page, req_id)
        else:
            logger.info(" UI状态已正确设置，无需更新")
            return True

    except asyncio.CancelledError:
        raise
    except Exception as e:
        logger.error(f" 验证并应用UI状态设置时发生错误: {e}")
        return False


async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) -> bool:
    """切换AI Studio模型"""
    logger.info(f" 开始切换模型到: {model_id}")
    original_prefs_str: Optional[str] = None
    original_prompt_model: Optional[str] = None
    new_chat_url = f"https://{AI_STUDIO_URL_PATTERN}prompts/new_chat"

    try:
        original_prefs_str = await page.evaluate(
            "() => localStorage.getItem('aiStudioUserPreference')"
        )
        if original_prefs_str:
            try:
                original_prefs_obj = json.loads(original_prefs_str)
                original_prompt_model = original_prefs_obj.get("promptModel")
                logger.info(
                    f" 切换前 localStorage.promptModel 为: {original_prompt_model or '未设置'}"
                )
            except json.JSONDecodeError:
                logger.warning(" 无法解析原始的 aiStudioUserPreference JSON 字符串。")
                original_prefs_str = None

        current_prefs_for_modification = (
            json.loads(original_prefs_str) if original_prefs_str else {}
        )
        full_model_path = f"models/{model_id}"

        if current_prefs_for_modification.get("promptModel") == full_model_path:
            logger.info(
                f" 模型已经设置为 {model_id} (localStorage 中已是目标值)，无需切换"
            )
            if page.url != new_chat_url:
                logger.info(
                    f" 当前 URL 不是 new_chat ({page.url})，导航到 {new_chat_url}"
                )
                await page.goto(
                    new_chat_url, wait_until="domcontentloaded", timeout=30000
                )
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                    timeout=30000
                )
            return True

        logger.info(
            f" 从 {current_prefs_for_modification.get('promptModel', '未知')} 更新 localStorage.promptModel 为 {full_model_path}"
        )
        current_prefs_for_modification["promptModel"] = full_model_path
        await page.evaluate(
            "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
            json.dumps(current_prefs_for_modification),
        )

        # 使用新的强制设置功能
        logger.info(" 应用强制UI状态设置...")
        ui_state_success = await _verify_and_apply_ui_state(page, req_id)
        if not ui_state_success:
            logger.warning(" UI状态设置失败，但继续执行模型切换流程")

        # 为了保持兼容性，也更新当前的prefs对象
        current_prefs_for_modification["isAdvancedOpen"] = True
        current_prefs_for_modification["areToolsOpen"] = True
        await page.evaluate(
            "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
            json.dumps(current_prefs_for_modification),
        )

        logger.info(f" localStorage 已更新，导航到 '{new_chat_url}' 应用新模型...")
        await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=30000)

        input_field = page.locator(INPUT_SELECTOR)
        await expect_async(input_field).to_be_visible(timeout=30000)
        logger.info(" 页面已导航到新聊天并加载完成，输入框可见")

        # 页面加载后再次验证UI状态设置
        logger.info(" 页面加载完成，验证UI状态设置...")
        final_ui_state_success = await _verify_and_apply_ui_state(page, req_id)
        if final_ui_state_success:
            logger.info(" UI状态最终验证成功")
        else:
            logger.warning(" UI状态最终验证失败，但继续执行模型切换流程")

        final_prefs_str = await page.evaluate(
            "() => localStorage.getItem('aiStudioUserPreference')"
        )
        final_prompt_model_in_storage: Optional[str] = None
        if final_prefs_str:
            try:
                final_prefs_obj = json.loads(final_prefs_str)
                final_prompt_model_in_storage = final_prefs_obj.get("promptModel")
            except json.JSONDecodeError:
                logger.warning(" 无法解析刷新后的 aiStudioUserPreference JSON 字符串。")

        if final_prompt_model_in_storage == full_model_path:
            logger.info(
                f" AI Studio localStorage 中模型已成功设置为: {full_model_path}"
            )

            page_display_match = False

            # 获取parsed_model_list
            from api_utils.server_state import state

            parsed_model_list = getattr(state, "parsed_model_list", [])

            if parsed_model_list:
                for m_obj in parsed_model_list:
                    if m_obj.get("id") == model_id:
                        m_obj.get("display_name")
                        break

            try:
                model_name_locator = page.locator('[data-test-id="model-name"]')
                actual_displayed_model_id_on_page_raw = (
                    await model_name_locator.first.inner_text(timeout=5000)
                )
                actual_displayed_model_id_on_page = (
                    actual_displayed_model_id_on_page_raw.strip()
                )

                target_model_id = model_id

                if actual_displayed_model_id_on_page == target_model_id:
                    page_display_match = True
                    logger.info(
                        f" 页面显示模型ID ('{actual_displayed_model_id_on_page}') 与期望ID ('{target_model_id}') 一致。"
                    )
                else:
                    page_display_match = False
                    logger.error(
                        f" 页面显示模型ID ('{actual_displayed_model_id_on_page}') 与期望ID ('{target_model_id}') 不一致。"
                    )

            except asyncio.CancelledError:
                raise
            except Exception as e_disp:
                page_display_match = False  # 读取失败则认为不匹配
                logger.warning(
                    f" 读取页面显示的当前模型ID时出错: {e_disp}。将无法验证页面显示。"
                )

            if page_display_match:
                try:
                    logger.info(" 模型切换成功，重新启用 '临时聊天' 模式...")
                    incognito_button_locator = page.locator(
                        'button[aria-label="Temporary chat toggle"]'
                    )

                    await incognito_button_locator.wait_for(
                        state="visible", timeout=5000
                    )

                    button_classes = await incognito_button_locator.get_attribute(
                        "class"
                    )

                    if button_classes and "ms-button-active" in button_classes:
                        logger.info(" '临时聊天' 模式已处于激活状态。")
                    else:
                        logger.info(" '临时聊天' 模式未激活，正在点击以开启...")
                        await incognito_button_locator.click(timeout=3000)
                        await asyncio.sleep(0.5)

                        updated_classes = await incognito_button_locator.get_attribute(
                            "class"
                        )
                        if updated_classes and "ms-button-active" in updated_classes:
                            logger.info(" '临时聊天' 模式已成功重新启用。")
                        else:
                            logger.warning(
                                " 点击后 '临时聊天' 模式状态验证失败，可能未成功重新开启。"
                            )

                except asyncio.CancelledError:
                    raise
                except Exception as e:
                    logger.warning(f" 模型切换后重新启用 '临时聊天' 模式失败: {e}")
                return True
            else:
                logger.error(
                    " 模型切换失败，因为页面显示的模型与期望不符 (即使localStorage可能已更改)。"
                )
        else:
            logger.error(
                f" AI Studio 未接受模型更改 (localStorage)。期望='{full_model_path}', 实际='{final_prompt_model_in_storage or '未设置或无效'}'."
            )

        logger.info(" 模型切换失败。尝试恢复到页面当前实际显示的模型的状态...")
        current_displayed_name_for_revert_raw = "无法读取"
        current_displayed_name_for_revert_stripped = "无法读取"

        try:
            model_name_locator_revert = page.locator('[data-test-id="model-name"]')
            current_displayed_name_for_revert_raw = (
                await model_name_locator_revert.first.inner_text(timeout=5000)
            )
            current_displayed_name_for_revert_stripped = (
                current_displayed_name_for_revert_raw.strip()
            )
            logger.info(
                f" 恢复：页面当前显示的模型名称 (原始: '{current_displayed_name_for_revert_raw}', 清理后: '{current_displayed_name_for_revert_stripped}')"
            )
        except asyncio.CancelledError:
            raise
        except Exception as e_read_disp_revert:
            logger.warning(
                f" 恢复：读取页面当前显示模型名称失败: {e_read_disp_revert}。将尝试回退到原始localStorage。"
            )
            if original_prefs_str:
                logger.info(
                    f" 恢复：由于无法读取当前页面显示，尝试将 localStorage 恢复到原始状态: '{original_prompt_model or '未设置'}'"
                )
                await page.evaluate(
                    "(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)",
                    original_prefs_str,
                )
                logger.info(
                    f" 恢复：导航到 '{new_chat_url}' 以应用恢复的原始 localStorage 设置..."
                )
                await page.goto(
                    new_chat_url, wait_until="domcontentloaded", timeout=20000
                )
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                    timeout=20000
                )
                logger.info(
                    " 恢复：页面已导航到新聊天并加载，已尝试应用原始 localStorage。"
                )
            else:
                logger.warning(
                    " 恢复：无有效的原始 localStorage 状态可恢复，也无法读取当前页面显示。"
                )
            return False

        model_id_to_revert_to = None
        if current_displayed_name_for_revert_stripped != "无法读取":
            model_id_to_revert_to = current_displayed_name_for_revert_stripped
            logger.info(
                f" 恢复：页面当前显示的ID是 '{model_id_to_revert_to}'，将直接用于恢复。"
            )
        else:
            if current_displayed_name_for_revert_stripped == "无法读取":
                logger.warning(
                    " 恢复：因无法读取页面显示名称，故不能从 parsed_model_list 转换ID。"
                )
            else:
                logger.warning(
                    f" 恢复：parsed_model_list 为空，无法从显示名称 '{current_displayed_name_for_revert_stripped}' 转换模型ID。"
                )

        if model_id_to_revert_to:
            base_prefs_for_final_revert = {}
            try:
                current_ls_content_str = await page.evaluate(
                    "() => localStorage.getItem('aiStudioUserPreference')"
                )
                if current_ls_content_str:
                    base_prefs_for_final_revert = json.loads(current_ls_content_str)
                elif original_prefs_str:
                    base_prefs_for_final_revert = json.loads(original_prefs_str)
            except json.JSONDecodeError:
                logger.warning(" 恢复：解析现有 localStorage 以构建恢复偏好失败。")

            path_to_revert_to = f"models/{model_id_to_revert_to}"
            base_prefs_for_final_revert["promptModel"] = path_to_revert_to
            # 使用新的强制设置功能
            logger.info(" 恢复：应用强制UI状态设置...")
            ui_state_success = await _verify_and_apply_ui_state(page, req_id)
            if not ui_state_success:
                logger.warning(" 恢复：UI状态设置失败，但继续执行恢复流程")

            # 为了保持兼容性，也更新当前的prefs对象
            base_prefs_for_final_revert["isAdvancedOpen"] = True
            base_prefs_for_final_revert["areToolsOpen"] = True
            logger.info(
                f" 恢复：准备将 localStorage.promptModel 设置回页面实际显示的模型的路径: '{path_to_revert_to}'，并强制设置配置选项"
            )
            await page.evaluate(
                "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
                json.dumps(base_prefs_for_final_revert),
            )
            logger.info(
                f" 恢复：导航到 '{new_chat_url}' 以应用恢复到 '{model_id_to_revert_to}' 的 localStorage 设置..."
            )
            await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=30000)
            await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                timeout=30000
            )

            # 恢复后再次验证UI状态
            logger.info(" 恢复：页面加载完成，验证UI状态设置...")
            final_ui_state_success = await _verify_and_apply_ui_state(page, req_id)
            if final_ui_state_success:
                logger.info(" 恢复：UI状态最终验证成功")
            else:
                logger.warning(" 恢复：UI状态最终验证失败")

            logger.info(
                f" 恢复：页面已导航到新聊天并加载。localStorage 应已设置为反映模型 '{model_id_to_revert_to}'。"
            )
        else:
            logger.error(
                f" 恢复：无法将模型恢复到页面显示的状态，因为未能从显示名称 '{current_displayed_name_for_revert_stripped}' 确定有效模型ID。"
            )
            if original_prefs_str:
                logger.warning(
                    f" 恢复：作为最终后备，尝试恢复到原始 localStorage: '{original_prompt_model or '未设置'}'"
                )
                await page.evaluate(
                    "(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)",
                    original_prefs_str,
                )
                logger.info(
                    f" 恢复：导航到 '{new_chat_url}' 以应用最终后备的原始 localStorage。"
                )
                await page.goto(
                    new_chat_url, wait_until="domcontentloaded", timeout=20000
                )
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                    timeout=20000
                )
                logger.info(
                    " 恢复：页面已导航到新聊天并加载，已应用最终后备的原始 localStorage。"
                )
            else:
                logger.warning(" 恢复：无有效的原始 localStorage 状态可作为最终后备。")

        return False

    except asyncio.CancelledError:
        raise
    except Exception:
        logger.exception(" 切换模型过程中发生严重错误")
        # 导入save_error_snapshot函数
        from .operations import save_error_snapshot

        await save_error_snapshot(f"model_switch_error_{req_id}")
        try:
            if original_prefs_str:
                logger.info(
                    f" 发生异常，尝试恢复 localStorage 至: {original_prompt_model or '未设置'}"
                )
                await page.evaluate(
                    "(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)",
                    original_prefs_str,
                )
                logger.info(
                    f" 异常恢复：导航到 '{new_chat_url}' 以应用恢复的 localStorage。"
                )
                await page.goto(
                    new_chat_url, wait_until="domcontentloaded", timeout=15000
                )
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                    timeout=15000
                )
        except asyncio.CancelledError:
            raise
        except Exception as recovery_err:
            logger.error(f" 异常后恢复 localStorage 失败: {recovery_err}")
        return False


def load_excluded_models(filename: str):
    """加载排除的模型列表"""
    from api_utils.server_state import state

    excluded_model_ids = getattr(state, "excluded_model_ids", set())

    excluded_file_path = os.path.join(os.path.dirname(__file__), "..", filename)
    try:
        if os.path.exists(excluded_file_path):
            with open(excluded_file_path, "r", encoding="utf-8") as f:
                loaded_ids = {line.strip() for line in f if line.strip()}
            if loaded_ids:
                excluded_model_ids.update(loaded_ids)
                state.excluded_model_ids = excluded_model_ids
                logger.info(
                    f"从 '{filename}' 加载了 {len(loaded_ids)} 个模型到排除列表: {excluded_model_ids}"
                )
            else:
                logger.info(
                    f"'{filename}' 文件为空或不包含有效的模型 ID，排除列表未更改。"
                )
        else:
            logger.info(f"模型排除列表文件 '{filename}' 未找到，排除列表为空。")
    except Exception as e:
        logger.error(f"从 '{filename}' 加载排除模型列表时出错: {e}", exc_info=True)


async def _handle_initial_model_state_and_storage(page: AsyncPage):
    """处理初始模型状态和存储"""
    from api_utils.server_state import state

    getattr(state, "current_ai_studio_model_id", None)
    getattr(state, "parsed_model_list", [])
    getattr(state, "model_list_fetch_event", None)

    logger.info("--- (新) 处理初始模型状态, localStorage 和 isAdvancedOpen ---")
    needs_reload_and_storage_update = False
    reason_for_reload = ""

    try:
        initial_prefs_str = await page.evaluate(
            "() => localStorage.getItem('aiStudioUserPreference')"
        )
        if not initial_prefs_str:
            needs_reload_and_storage_update = True
            reason_for_reload = "localStorage.aiStudioUserPreference 未找到。"
            logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
        else:
            logger.info("   localStorage 中找到 'aiStudioUserPreference'。正在解析...")
            try:
                pref_obj = json.loads(initial_prefs_str)
                prompt_model_path = pref_obj.get("promptModel")
                pref_obj.get("isAdvancedOpen")
                is_prompt_model_valid = (
                    isinstance(prompt_model_path, str) and prompt_model_path.strip()
                )

                if not is_prompt_model_valid:
                    needs_reload_and_storage_update = True
                    reason_for_reload = "localStorage.promptModel 无效或未设置。"
                    logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
                else:
                    # 使用新的UI状态验证功能
                    ui_state = await _verify_ui_state_settings(page, "initial")
                    if ui_state["needsUpdate"]:
                        needs_reload_and_storage_update = True
                        reason_for_reload = f"UI状态需要更新: isAdvancedOpen={ui_state['isAdvancedOpen']}, areToolsOpen={ui_state['areToolsOpen']} (期望: True)"
                        logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
                    else:
                        state.current_ai_studio_model_id = prompt_model_path.split("/")[
                            -1
                        ]
                        logger.info(
                            f"   localStorage 有效且UI状态正确。初始模型 ID 从 localStorage 设置为: {state.current_ai_studio_model_id}"
                        )
            except json.JSONDecodeError:
                needs_reload_and_storage_update = True
                reason_for_reload = (
                    "解析 localStorage.aiStudioUserPreference JSON 失败。"
                )
                logger.error(f"   判定需要刷新和存储更新: {reason_for_reload}")

        if needs_reload_and_storage_update:
            logger.info(f"   执行刷新和存储更新流程，原因: {reason_for_reload}")
            logger.info(
                "   步骤 1: 调用 _set_model_from_page_display(set_storage=True) 更新 localStorage 和全局模型 ID..."
            )
            await _set_model_from_page_display(page, set_storage=True)

            current_page_url = page.url
            logger.info(
                f"   步骤 2: 重新加载页面 ({current_page_url}) 以应用 isAdvancedOpen=true..."
            )
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    logger.info(
                        f"   尝试重新加载页面 (第 {attempt + 1}/{max_retries} 次): {current_page_url}"
                    )
                    await page.goto(
                        current_page_url, wait_until="domcontentloaded", timeout=40000
                    )
                    await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                        timeout=30000
                    )
                    logger.info(f"   页面已成功重新加载到: {page.url}")

                    # 页面重新加载后验证UI状态
                    logger.info("   页面重新加载完成，验证UI状态设置...")
                    reload_ui_state_success = await _verify_and_apply_ui_state(
                        page, "reload"
                    )
                    if reload_ui_state_success:
                        logger.info("   重新加载后UI状态验证成功")
                    else:
                        logger.warning("   重新加载后UI状态验证失败")

                    break  # 成功则跳出循环
                except asyncio.CancelledError:
                    raise
                except Exception as reload_err:
                    logger.warning(
                        f"   页面重新加载尝试 {attempt + 1}/{max_retries} 失败: {reload_err}"
                    )
                    if attempt < max_retries - 1:
                        logger.info("   将在5秒后重试...")
                        await asyncio.sleep(5)
                    else:
                        logger.error(
                            f"   页面重新加载在 {max_retries} 次尝试后最终失败: {reload_err}. 后续模型状态可能不准确。",
                            exc_info=True,
                        )
                        from .operations import save_error_snapshot

                        await save_error_snapshot(
                            f"initial_storage_reload_fail_attempt_{attempt + 1}"
                        )

            logger.info(
                "   步骤 3: 重新加载后，再次调用 _set_model_from_page_display(set_storage=False) 以同步全局模型 ID..."
            )
            await _set_model_from_page_display(page, set_storage=False)
            logger.info(
                f"   刷新和存储更新流程完成。最终全局模型 ID: {state.current_ai_studio_model_id}"
            )
        else:
            logger.info(
                "   localStorage 状态良好 (isAdvancedOpen=true, promptModel有效)，无需刷新页面。"
            )
    except asyncio.CancelledError:
        raise
    except Exception as e:
        logger.error(
            f"(新) 处理初始模型状态和 localStorage 时发生严重错误: {e}",
            exc_info=True,
        )
        try:
            logger.warning(
                "   由于发生错误，尝试回退仅从页面显示设置全局模型 ID (不写入localStorage)..."
            )
            await _set_model_from_page_display(page, set_storage=False)
        except asyncio.CancelledError:
            raise
        except Exception as fallback_err:
            logger.error(f"   回退设置模型ID也失败: {fallback_err}")


async def _set_model_from_page_display(page: AsyncPage, set_storage: bool = False):
    """从页面显示设置模型"""
    from api_utils.server_state import state

    getattr(state, "current_ai_studio_model_id", None)
    getattr(state, "parsed_model_list", [])
    model_list_fetch_event = getattr(state, "model_list_fetch_event", None)

    try:
        logger.info("   尝试从页面显示元素读取当前模型名称...")
        model_name_locator = page.locator('[data-test-id="model-name"]')
        displayed_model_name_from_page_raw = await model_name_locator.first.inner_text(
            timeout=7000
        )
        displayed_model_name = displayed_model_name_from_page_raw.strip()
        logger.info(
            f"   页面当前显示模型名称 (原始: '{displayed_model_name_from_page_raw}', 清理后: '{displayed_model_name}')"
        )

        found_model_id_from_display = None
        if model_list_fetch_event and not model_list_fetch_event.is_set():
            logger.info("   等待模型列表数据 (最多5秒) 以便转换显示名称...")
            try:
                await asyncio.wait_for(model_list_fetch_event.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                logger.warning("   等待模型列表超时，可能无法准确转换显示名称为ID。")

        found_model_id_from_display = displayed_model_name
        logger.info(f"   页面显示的直接是模型ID: '{found_model_id_from_display}'")

        new_model_value = found_model_id_from_display
        if state.current_ai_studio_model_id != new_model_value:
            state.current_ai_studio_model_id = new_model_value
            logger.info(
                f"   全局 current_ai_studio_model_id 已更新为: {state.current_ai_studio_model_id}"
            )
        else:
            logger.info(
                f"   全局 current_ai_studio_model_id ('{state.current_ai_studio_model_id}') 与从页面获取的值一致，未更改。"
            )

        if set_storage:
            logger.info(
                "   准备为页面状态设置 localStorage (确保 isAdvancedOpen=true)..."
            )
            existing_prefs_for_update_str = await page.evaluate(
                "() => localStorage.getItem('aiStudioUserPreference')"
            )
            prefs_to_set = {}
            if existing_prefs_for_update_str:
                try:
                    prefs_to_set = json.loads(existing_prefs_for_update_str)
                except json.JSONDecodeError:
                    logger.warning(
                        "   解析现有 localStorage.aiStudioUserPreference 失败，将创建新的偏好设置。"
                    )

            # 使用新的强制设置功能
            logger.info("     应用强制UI状态设置...")
            ui_state_success = await _verify_and_apply_ui_state(page, "set_model")
            if not ui_state_success:
                logger.warning("     UI状态设置失败，使用传统方法")
                prefs_to_set["isAdvancedOpen"] = True
                prefs_to_set["areToolsOpen"] = True
            else:
                # 确保prefs_to_set也包含正确的设置
                prefs_to_set["isAdvancedOpen"] = True
                prefs_to_set["areToolsOpen"] = True
            logger.info("     强制 isAdvancedOpen: true, areToolsOpen: true")

            if found_model_id_from_display:
                new_prompt_model_path = f"models/{found_model_id_from_display}"
                prefs_to_set["promptModel"] = new_prompt_model_path
                logger.info(
                    f"     设置 promptModel 为: {new_prompt_model_path} (基于找到的ID)"
                )
            elif "promptModel" not in prefs_to_set:
                logger.warning(
                    f"     无法从页面显示 '{displayed_model_name}' 找到模型ID，且 localStorage 中无现有 promptModel。promptModel 将不会被主动设置以避免潜在问题。"
                )

            default_keys_if_missing = {
                "bidiModel": "models/gemini-1.0-pro-001",
                "isSafetySettingsOpen": False,
                "hasShownSearchGroundingTos": False,
                "autosaveEnabled": True,
                "theme": "system",
                "bidiOutputFormat": 3,
                "isSystemInstructionsOpen": False,
                "warmWelcomeDisplayed": True,
                "getCodeLanguage": "Node.js",
                "getCodeHistoryToggle": False,
                "fileCopyrightAcknowledged": True,
            }
            for key, val_default in default_keys_if_missing.items():
                if key not in prefs_to_set:
                    prefs_to_set[key] = val_default

            await page.evaluate(
                "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
                json.dumps(prefs_to_set),
            )
            logger.info(
                f"   localStorage.aiStudioUserPreference 已更新。isAdvancedOpen: {prefs_to_set.get('isAdvancedOpen')}, areToolsOpen: {prefs_to_set.get('areToolsOpen')} (期望: True), promptModel: '{prefs_to_set.get('promptModel', '未设置/保留原样')}'。"
            )
    except asyncio.CancelledError:
        raise
    except Exception as e_set_disp:
        logger.error(f"   尝试从页面显示设置模型时出错: {e_set_disp}", exc_info=True)
